import cv2
import socket
import pickle
import struct
import time
import argparse
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class VideoStreamClient:
    def __init__(self, host='localhost', port=5000):
        self.host = host
        self.port = port
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.running = False
        self.last_heartbeat = time.time()
        self.heartbeat_interval = 5  # 心跳间隔(秒)
        self.reconnect_attempts = 0
        self.max_reconnect_attempts = 5
        
    def connect(self):
        try:
            self.client_socket.connect((self.host, self.port))
            self.running = True
            logger.info(f"已连接到服务器 {self.host}:{self.port}")
            return True
        except Exception as e:
            logger.error(f"连接失败: {e}")
            return False
    
    def send_heartbeat(self):
        try:
            self.client_socket.sendall(struct.pack("L", 0))  # 发送0作为心跳包
            return True
        except:
            return False
            
    def check_connection(self):
        current_time = time.time()
        if current_time - self.last_heartbeat > self.heartbeat_interval:
            if not self.send_heartbeat():
                logger.warning("心跳失败，尝试重新连接...")
                return self.reconnect()
            self.last_heartbeat = current_time
        return True
            
    def reconnect(self):
        self.client_socket.close()
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.reconnect_attempts += 1
        
        if self.reconnect_attempts > self.max_reconnect_attempts:
            logger.error("达到最大重连次数")
            self.running = False
            return False
            
        try:
            self.client_socket.connect((self.host, self.port))
            self.reconnect_attempts = 0
            logger.info("重新连接成功")
            return True
        except Exception as e:
            logger.error(f"重新连接失败: {e}")
            time.sleep(1)  # 等待1秒后重试
            return self.reconnect()
            
    def send_video(self, camera_index=0):
        cap = cv2.VideoCapture(camera_index)
        if not cap.isOpened():
            logger.error("无法打开视频源")
            return
            
        try:
            while self.running:
                if not self.check_connection():
                    break
                    
                ret, frame = cap.read()
                if not ret:
                    break
                    
                # 压缩帧为JPEG格式
                _, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
                
                # 序列化帧数据
                data = pickle.dumps(buffer)
                message_size = struct.pack("L", len(data))
                
                # 发送帧数据(带流量控制)
                try:
                    # 检查发送缓冲区(避免堆积)
                    if self.client_socket.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF) < len(data) * 2:
                        time.sleep(0.05)  # 缓冲区快满时稍作等待
                        
                    self.client_socket.sendall(message_size + data)
                except socket.error as e:
                    logger.error(f"发送错误: {e}")
                    if not self.reconnect():
                        break
                        
                time.sleep(0.03)  # 控制帧率
                
        finally:
            cap.release()
            self.client_socket.close()
            logger.info("视频流传输已停止")
    
    def start(self, camera_index=0):
        if self.connect():
            self.send_video(camera_index)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Video Streaming Client')
    parser.add_argument('--host', default='localhost', help='Server host address')
    parser.add_argument('--port', type=int, default=5000, help='Server port number')
    parser.add_argument('--camera', type=int, default=0, help='Camera index')
    
    args = parser.parse_args()
    
    client = VideoStreamClient(args.host, args.port)
    client.start(args.camera)
